# Author：gyx
# Usage: build uav with robot arm dynamics model
# Last change: 2024-7-20

from common_function import *
from init_param import UAVParameters
import numpy as np


class UAVModel:
    def __init__(self, params=UAVParameters()):
        self.params = params
        self.w_current = params.ModelInit_Rads
        self.rateB_dot = params.ModelInit_RateB
        # self.w_manipulator = 0
        # self.xi = np.concatenate(params.ModelInit_PosE, params.ModelInit_AngEuler, params.ModelInit_Jointq)
        self.I1 = np.array([[0, 0, 0],
                            [0, 0, 0],
                            [0, 0, params.linkinertia1]])
        self.I2 = np.array([[0, 0, 0],
                            [0, 0, 0],
                            [0, 0, params.linkinertia2]])
        self.r = np.array([0.05, 0.075])
        self.l = np.array([0.1, 0.15])

    def compute_r_out(self, theta1):
        r_out = np.array([self.r[0], 0, 0]) + self.l[0] / 2 * np.array([np.cos(theta1), np.sin(theta1), 0])
        return r_out

    def get_angular_velocity_derivative(self, M, J, omega):
        return np.linalg.inv(J) @ (M - np.cross(omega, J @ omega))

    def get_linear_velocity_derivative(self, F, m):
        return np.clip(F / m, -0.4, 0.4)

    def Obtain_force_torque(self, w, R, Cm, Ct, Vb, Cd, wb, Cdm, Jrp):
        M_rctcm = np.array([
            [-np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
             -np.sin(np.pi / 4) * R * Ct],  # Roll torque
            [np.sin(np.pi / 4) * R * Ct, -np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
             -np.sin(np.pi / 4) * R * Ct],  # Pitch torque
            [Cm, Cm, -Cm, -Cm]  # Yaw torque
        ])

        Mp = np.dot(M_rctcm, w ** 2)  # Torque

        Fp = np.array([0, 0, -np.sum(Ct * (w ** 2))])  # Thrust
        Fd = -Cd * Vb * np.abs(Vb) * 0.5  # Aerodynamic force
        Md = -Cdm * wb * np.abs(wb)  # Aerodynamic moment

        # Gyro moment
        Ga = np.zeros(3)
        Ga[0] = Jrp * wb[1] * (w[0] + w[1] - w[2] - w[3])
        Ga[1] = Jrp * wb[0] * (-w[0] - w[1] + w[2] + w[3])

        return Fp, Fd, Mp, Md, Ga

    def motor_dynamics(self, w_current, w_target, T, dt):
        dw = (w_target - w_current) / T
        return w_current + dw * dt

    def manipulator_dynamics(self, state_man, omega_state, a, R_i_b):
        # velB, angEuler, rateB = state
        omega, omega_dot = omega_state
        theta_m, theta_m_dot, theta_m_ddot = state_man
        R_b_0 = rotation_matrix([0, -np.pi / 2, np.pi / 2]).T
        R_1_0 = rot_dh(theta_m[0], 0)
        R_2_1 = rot_dh(theta_m[1], 0)
        R_b_1 = R_1_0.T @ R_b_0
        R_2_0 = R_1_0 @ R_2_1
        a_b = R_i_b @ a
        b1 = R_1_0.T
        # print(b1 @ np.array([[0], [0], theta_m_dot[0]]).reshape(3,1))
        omega_1 = R_b_1 @ omega + b1@np.array([0 ,0, theta_m_dot[0]])
        alpha_1 = R_b_1 @ omega_dot + b1@np.array([0, 0, theta_m_ddot[0]]) \
                  + cross_product(omega_1, b1@np.array([0,0, theta_m_dot[0]]))
        ac1 = R_b_1 @ a_b + cross_product(alpha_1, np.array([self.r[0], 0, 0])) \
              + cross_product(omega_1, cross_product(omega_1, np.array([self.r[0], 0, 0])))
        # print(omega)
        ae1 = R_b_1 @ a_b + cross_product(alpha_1, np.array([self.l[0], 0, 0])) \
              + cross_product(omega_1, cross_product(omega_1, np.array([self.l[0], 0, 0])))
        b2 = R_2_0.T
        omega_2 = R_2_1.T @ omega_1 + b2 @ np.array([0, 0, np.sum(theta_m_dot)])
        alpha_2 = R_b_1 @ omega_dot + b2 @ np.array([0, 0, np.sum(theta_m_ddot)]) \
                  + cross_product(omega_2, b2 @ np.array([0, 0, np.sum(theta_m_dot)]))

        ac2 = R_2_1.T @ ae1 + cross_product(alpha_2, np.array([self.r[1], 0, 0])) \
              + cross_product(omega_2, cross_product(omega_2, np.array([self.r[1], 0, 0])))
        g2 = R_2_0.T @ R_b_0 @ R_i_b @ np.array([0, 0, 9.8])
        g1 = R_1_0.T @ R_b_0 @ R_i_b @ np.array([0, 0, 9.8])
        f2 = self.params.linkmass2 * ac2 - self.params.linkmass2 * g2
        t2 = self.I2 @ alpha_2 -cross_product(f2, np.array([self.r[1], 0, 0]))\
             + cross_product(omega_2, self.I2 @ omega_2)

        self.f1 = R_2_1 @ f2 + self.params.linkmass1 * ac1 - self.params.linkmass1 * g1
        t1 = R_2_1 @ t2 - cross_product(self.f1, np.array([self.r[1], 0, 0])) \
             + cross_product(R_2_1 @ f2, np.array([self.l[0] - self.r[0], 0, 0])) \
             + self.I1 @ alpha_1 + cross_product(omega_1, self.I1 @ omega_1)
        # print(cross_product(R_2_1 @ f2, np.array([self.l[0] - self.r[0], 0, 0])))
        theta_m_ddot = [t1 / self.params.linkinertia1, t2 / self.params.linkinertia2]
        self.T_m = R_b_1.T @ t1
        self.F_m = R_b_1.T @ self.f1

        print('T_m',self.T_m,'F_m',self.F_m)
        # return [self.F_m, self.T_m]

    def uav_dynamics(self, state, w_target, theta1, dt):
        velB, angEuler, rateB = state

        m = self.params.ModelParam_uavMass
        J = self.params.ModelParam_uavJ
        g = self.params.ModelParam_envGravityAcc
        uavType = self.params.ModelParam_uavType
        R = self.params.ModelParam_uavR
        Cm = self.params.ModelParam_rotorCm
        Ct = self.params.ModelParam_rotorCt
        Cd = self.params.ModelParam_uavCd
        Cdm = self.params.ModelParam_uavCCm
        Jrp = self.params.ModelParam_motorJm
        T = self.params.ModelParam_motorT

        self.w_current = self.motor_dynamics(self.w_current, w_target, T, dt)
        Fp, Fd, Mp, Md, Ga = self.Obtain_force_torque(self.w_current, R, Cm, Ct, velB, Cd, rateB, Cdm, Jrp)
        dcm = dcm_from_euler(angEuler)
        total_force = Fp + self.F_m + Fd + np.array([0, 0, m * g]).dot(dcm)
        r_out = self.compute_r_out(theta1)
        total_torque = Mp + Md + Ga + cross_product(r_out, self.f1) + self.T_m
        velB_dot = self.get_linear_velocity_derivative(total_force, m)
        self.rateB_dot = self.get_angular_velocity_derivative(total_torque, J, rateB)
        return velB_dot, self.rateB_dot

    def update_state(self, state, state_man, theta1, w_target, dt):
        velB, angEuler, rateB = state
        R_i_b = rotation_matrix(angEuler).T
        # Runge-Kutta 4th Order Method
        k1_omega_state = [rateB , self.rateB_dot]
        self.manipulator_dynamics(state_man, k1_omega_state, velB , R_i_b)
        k1_velB_dot, k1_rateB_dot = self.uav_dynamics(state, w_target, theta1, dt)

        k2_omega_state = [rateB, k1_rateB_dot]
        # 更新机械臂动力学
        self.manipulator_dynamics(state_man, k2_omega_state, velB + 0.5 * k1_velB_dot, R_i_b)
        k2_velB_dot, k2_rateB_dot = self.uav_dynamics(
            (velB + 0.5 * k1_velB_dot * dt, angEuler, rateB + 0.5 * k1_rateB_dot * dt), w_target, theta1, dt)

        k3_omega_state = [rateB + 0.5 * k1_rateB_dot, k2_rateB_dot]
        self.manipulator_dynamics(state_man, k3_omega_state, velB + 0.5 * k2_velB_dot * dt, R_i_b)
        k3_velB_dot, k3_rateB_dot = self.uav_dynamics(
            (velB + 0.5 * k2_velB_dot * dt, angEuler, rateB + 0.5 * k2_rateB_dot * dt), w_target, theta1, dt)

        k4_omega_state = [rateB + k3_rateB_dot, k3_rateB_dot]
        self.manipulator_dynamics(state_man, k4_omega_state, velB +  k3_velB_dot * dt, R_i_b)
        k4_velB_dot, k4_rateB_dot = self.uav_dynamics(
            (velB + k3_velB_dot * dt, angEuler, rateB + k3_rateB_dot * dt), w_target, theta1, dt)


        velB = velB + (k1_velB_dot + 2 * k2_velB_dot + 2 * k3_velB_dot + k4_velB_dot) * dt / 6
        rateB = rateB + (k1_rateB_dot + 2 * k2_rateB_dot + 2 * k3_rateB_dot + k4_rateB_dot) * dt / 6
        return velB, angEuler, rateB
